import scipy.io as sio
import numpy as np
import pywt
import matplotlib.pyplot as plt
from scipy.signal import butter, lfilter


# 对手环获取信号进行折叠，让它们滚去该去的通道
def channel_reconstruct(signal, data_points):
    channeled_data = {'emg': [], 'label': []}
    step = 8
    num_batch = data_points // step
    remainder = data_points % step
    signal_shape = (num_batch, step)
    label_shape = (num_batch, 1)
    new_signal = np.zeros(signal_shape)
    new_label = np.zeros(label_shape)
    # 初始化新数据and标签

    start_index = 0
    for i in range(num_batch):
        end_index = start_index + step
        channel_data = signal['emg'][start_index:end_index, :]
        new_signal[i] = channel_data
        start_index = end_index
    if remainder > 0:
        channel_data = signal['emg'][-start_index, :]
        padding = np.zeros(step)
        channel_data = padding
        new_signal[-1] = channel_data
    start_index_1 = 0
    for i in range(num_batch):
        new_label[i] = signal['label'][start_index, :]
        end_index_1 = start_index_1 + step
        start_index_1 = end_index_1

    channeled_data["emg"] = new_signal
    channeled_data["label"] = new_label
    print('channel_reconstruct')
    print(f'emg', channeled_data['emg'].shape)
    print(f'label', channeled_data['label'].shape)
    print('---------------------------------------------')
    return channeled_data


# 华东窗口（linux输入法真的很灾难）
def data_process_with_sliding_window(signal, label_original, window_size, step_size, method):
    spectrograms_windowed = []
    start = 0
    # 对数据进行小波变换
    for start in range(0, len(signal) - window_size + 1, step_size):
        print(f'start of window:', start)
        end = start + window_size
        windowed_signal = signal[start:end]
        # print(f'windowed shape',windowed_signal.shape)
        # plt.figure(figsize=(16, 12))
        # plt.subplot(224)
        # plt.plot(windowed_signal,label=f'm3 bro')
        # plt.legend()
        # plt.title('Wavelet Transform Spectrogram old')
        # plt.xlabel('Time [sec]')
        # plt.ylabel('Amplitude')

        if method == 'wavelet':  # 先使用处理后获得的逼近系数来进行实验*
            coeffs = pywt.wavedec(windowed_signal, 'db2', level=2)
            cA_2, cD2, cD1 = coeffs
            # x_recon = pywt.waverec(coeffs, 'db2')
            # print(f'coeffs',coeffs)
            # print(f'ca-2',cA_2.shape)
            # print(f'cd2',cD2.shape)
            # print(f'cd1',cD1.shape)
            # 绘制每一级的逼近系数
            # for i, c in enumerate(coeffs):
            #     plt.subplot(2, 2, i+1)  # 2 * i + 1 为奇数位置的子图
            #     plt.plot(c, label=f'Approximation coefficients (Level {i})')
            #     plt.legend()
            #
            # plt.tight_layout()
            # plt.show()

            spectrograms_windowed.extend(cA_2)
    # 对小波变换后的数据再加上正确的标签，处理前的信号长度如果是1510的话，50个窗口生成2000个数据
    # 总之如果我们在采集时让停顿和动作都为3s的话刚好基本可以达到这个效果
    # 以此思路对标签进行再处理,取这个窗口中间的标签作为这个窗口的标签，而且要将数据分为（batch_size,channel,length）
    spectrograms_label = []
    for start in range(0, len(label_original) - window_size + 1, step_size):
        print(f'window s label:', start)
        window_s_label = label_original[(start + window_size / 2)]
        spectrograms_label.extend(window_s_label)

    return spectrograms_windowed, spectrograms_label

#把滑动窗口处理过的数据按照动作进行拼接
def processed_data_join(spectrogram,window_label):
    integrated_datas = {'emg': [], 'label': []}
    data = np.array(spectrogram)
    data_l = np.array(window_label)
    print(data.shape)
    print(data_l.shape)
    num_data_points = data.size[0]
    length = 2000  # 这个得手动修改，代码能力卡这儿了
    num_batch = num_data_points // length
    remainder = num_data_points % length
    # 折叠后的数据维度
    spectrogram_shape = (num_batch, 4, length)
    label_shape = (num_batch, )
    new_spectrogram = np.zeros(spectrogram_shape)
    new_label = np.zeros(label_shape)

    # 开始拼接......！
    start_index = 0
    for i in range(num_batch):
        end_index = start_index + length
        batch_data = spectrogram[:, start_index:end_index]
        new_spectrogram[i] = batch_data
        start_index = end_index
    if remainder > 0:
        # 提取剩余的数据片段
        batch_data = spectrogram[:, -start_index: ]
        # 为了使其能够匹配新数据的形状，我们需要在前面填充零
        padding = np.zeros((length, 4))
        batch_data = padding
        # 将数据片段存储到新数据数组中
        new_spectrogram[-1] = batch_data
    start_index_1 = length/2
    for i in range(num_batch):
        new_label[i] = window_label[start_index_1 + length, :]
        end_index_1 = start_index_1 + length
        start_index_1 = end_index_1
    integrated_datas["emg"] = new_spectrogram
    integrated_datas["label"] = new_label

    print(f'emg:', integrated_datas['emg'].shape)
    print(f'label', integrated_datas['label'].shape)
    sio.savemat('ugodata-process-to-use.mat', integrated_datas)
    return integrated_datas
# 滤出50hz电源信号
def data_process_to_50hz(signal, fs):
    cutoff_freq = 50
    sample_rate = fs
    order = 2
    nyq = 0.5 * sample_rate
    normal_cutoff_up = cutoff_freq / nyq
    normal_cutoff_down = (cutoff_freq - 1) / nyq
    b, a = butter(order, [normal_cutoff_down, normal_cutoff_up], btype='bandstop', analog=False)
    # a = butter(order, normal_cutoff, btype='high', analog=False)
    filtered_signal = lfilter(b, a, signal)
    print(f'filtered', filtered_signal.shape)
    plt.figure(figsize=(12, 6))
    plt.subplot(211)
    # plt.plot(signal)
    # plt.title('Original Signal with 50Hz Noise')
    # plt.xlabel('Sample')
    # plt.ylabel('Amplitude')

    plt.subplot(211)
    plt.plot(filtered_signal)
    plt.title('Filtered Signal')
    plt.xlabel('Sample')
    plt.ylabel('Amplitude')

    return filtered_signal


# 开始运行
# 总之先把数据赛进来看看到底是怎么一回事
data = sio.loadmat('ugodata.mat')
emg = data['emg']
print(f'emgshape', emg.shape)
label = data['label']
print(f'labelshape', label.shape)

# 之后再用的东西，等8通道数据捕出来了再说
# data_preprocess = channel_reconstruct(data)

data_process_to_50hz(emg, 101)

spectrograms, label = data_process_with_sliding_window(emg, label, 40, 30, 'wavelet')

spec = np.array(spectrograms)

plt.subplot(212)
plt.plot(spec)
plt.title('with 50hz cutoff and wavelet')
plt.xlabel('shit')
plt.ylabel('Amplitude')
print(f'spec', spec.shape)
plt.tight_layout()
plt.show()

processed_data_join(spectrograms, label)

# 接下来对其进行处理，使用滑动窗口以及小波变换
